import streamlit as st
# 全新的AI助手，专用负责翻译
from langchain_openai import ChatOpenAI
from langchain import PromptTemplate
import base64
import data.data as dd
#ChapOpenAI提供了各种大模型工具
#调用大模型

def main_bg(main_bg):
    main_bg_ext = "jpg"
    st.markdown(
        f"""
         <style>
         .stApp {{
             background: url(data:image/{main_bg_ext};base64,{base64.b64encode(open(main_bg, "rb").read()).decode()});
             background-size: cover
         }}
         </style>
         """,
        unsafe_allow_html=True)
main_bg('login.jpg')
def create_model():
    llm= ChatOpenAI(
        model="glm-4-0520",
        api_key="449efb54f05261e4d7ed49c7d81d535f.PVBWdC8gisFE0a3t",
        temperature="0.1",
        base_url="https://open.bigmodel.cn/api/paas/v4/"
)
    return llm

def model_invoke(message):
    llm = create_model()
    result = llm.invoke(message)
    return result.content
#翻译机器人模型
def create_translation_prompt_chain():
    prop = PromptTemplate(
        input_variables=["context"],
        template="""
           你是一个专业的翻译学家，只能进行与翻译有关的工作，将中文翻译为英语，将英语翻译为中文，其他问题你都回答不知道。用户输入的问题是：{context}
        """
    )
    llm = create_model()
    chain = prop | llm
    return chain
# 通过langchain链实现结果的调用
def chain_invoke(message):
    chain = create_translation_prompt_chain()
    result = chain.invoke(message)
    return result.content
user_id = st.session_state.user_id  # 用户id就是某一个用户的唯一标识
username = st.session_state.username

# streamlit的session_state缓存器进行处理 新数据覆盖旧数据的问题




input = st.chat_input("请输入你想要翻译的内容")
if input:
    with st.chat_message("user"):
        st.write(input)
    dd.add_chat_message_from_el(user_id, input, "user")
    result = chain_invoke({"context":input})
    with st.chat_message("assistant"):
        st.write(result)
    dd.add_chat_message_from_el(user_id, result, "assistant")

col1,col2= st.columns([1,1])
with col1:
    bt3 = st.button("历史记录")
    if bt3:
        history = dd.query_message_by_user_id_from_el(user_id=user_id)
        if history:
            for msg in history:
                with st.chat_message(msg["role"]):
                    st.write(msg["message"])
        else:
            with st.chat_message("assistant"):
                st.write("你好，我是你的红绿蓝助手，可以回答你的任何的翻译问题")

with col2:
    bt3 = st.button("创建一个新的会话")
    if bt3:
        st.switch_page("pages/translation_chatbot_select.py")